"""
验证 Evaluator 判别能力的对照实验：
- 对每个疾病生成若干问答；
- 通过字符串替换构造“错误疾病”版本；
- 比较 evaluator 在原始 vs 扰动版本的打分差异。
"""

import random
import json
import os
from tqdm import tqdm
from typing import List, Dict, Any
from medInfoQa import QAAgentPipeline ,read_json 

# ============ 参数 ============
N_DISEASE_SAMPLE = 20
N_GENERATE_PER_DISEASE = 3
RANDOM_SEED = 42
OUT_PATH = "./experiments/eval_diagnosis_control_result.json"

random.seed(RANDOM_SEED)

# ============ 主逻辑 ============

def run_control_experiment(gen_model_id="pretrained/medgemma4b", eval_model_id="pretrained/medgemma4b",
                           disease_path="../fundus-reasoner-dataprocess/configs/all_cn_name.json"):

    # 初始化 pipeline
    pipe = QAAgentPipeline(gen_model_id, eval_model_id)

    # 读取疾病清单
    all_diseases: List[str] = read_json(disease_path)
    if len(all_diseases) > N_DISEASE_SAMPLE:
        diseases = random.sample(all_diseases, N_DISEASE_SAMPLE)
    else:
        diseases = all_diseases

    # 用于统计
    all_records = []
    tar_accept = 0  # True Accept Rate (原始正确率)
    trr_reject = 0  # True Reject Rate (扰动拒绝率)
    total = 0

    print(f"[INFO] 共选取 {len(diseases)} 个疾病进行评估对照实验...\n")

    for dis in tqdm(diseases, desc="评估中"):
        # Step 1: 生成问答
        parsed, _ = pipe.single_for_disease(disease_cn=dis, n_generate=N_GENERATE_PER_DISEASE, keep_top_k=N_GENERATE_PER_DISEASE)
        if not parsed:
            continue

        # Step 2: 选一个错误疾病进行替换
        wrong_disease = random.choice([x for x in all_diseases if x != dis])
        for qa in parsed:
            q, a = qa["question"], qa["answer"]

            # 生成扰动版本（简单替换疾病名称）
            pert_q = q.replace(dis, wrong_disease)
            pert_a = a.replace(dis, wrong_disease)

            # Step 3: 评估两组问答
            eval_true = pipe.eva.evaluate_qa(q, a, context=f"原始疾病={dis}")
            eval_pert = pipe.eva.evaluate_qa(pert_q, pert_a, context=f"扰动疾病={wrong_disease}")

            # 计算结果
            def _decision_to_num(e):
                if not e:
                    return 0.5
                d = (e.get("decision") or e.get("is_correct") or "").lower() if isinstance(e.get("decision") or e.get("is_correct"), str) else e.get("is_correct")
                conf = e.get("confidence", 0.5)
                if isinstance(conf, str):
                    try:
                        conf = float(conf.strip("%")) / (100 if "%" in conf else 1)
                    except Exception:
                        conf = 0.5
                base = 1.0 if d in ("accept", "true", True) else (0.0 if d in ("reject", "false", False) else 0.5)
                return base * max(0.0, min(1.0, float(conf)))

            score_true = _decision_to_num(eval_true)
            score_pert = _decision_to_num(eval_pert)

            all_records.append({
                "disease": dis,
                "wrong_disease": wrong_disease,
                "question": q,
                "answer": a,
                "eval_true": eval_true,
                "eval_pert": eval_pert,
                "score_true": score_true,
                "score_pert": score_pert,
                "delta": round(score_true - score_pert, 3)
            })

            total += 1
            if score_true > 0.6:  # 判定为正确
                tar_accept += 1
            if score_pert < 0.4:  # 扰动被拒绝
                trr_reject += 1

    # Step 4: 汇总指标
    tar = round(tar_accept / total, 3) if total else 0
    trr = round(trr_reject / total, 3) if total else 0
    selectivity = round(tar - trr, 3)

    summary = {
        "num_diseases": len(diseases),
        "num_samples": total,
        "true_accept_rate": tar,
        "true_reject_rate": trr,
        "selectivity": selectivity,
    }

    print("\n========== [评估结果] ==========")
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    print(f"结果明细已保存至: {OUT_PATH}")

    os.makedirs(os.path.dirname(OUT_PATH), exist_ok=True)
    with open(OUT_PATH, "w", encoding="utf-8") as f:
        json.dump({"summary": summary, "records": all_records}, f, ensure_ascii=False, indent=2)


if __name__ == "__main__":
    run_control_experiment()
